#! /usr/bin/env python
# -*- coding: utf-8 -*-
# vim:fenc=utf-8
#
# Copyright © 2019 crane <crane@gosun>
#
# Distributed under terms of the MIT license.

"""
https://www.lintcode.com/problem/subarray-sum/description
"""


class Solution:
    """
    @param nums: A list of integers
    @return: A list of integers includes the index of the first number and the index of the last number
    """
    def subarraySum(self, nums):
        self.nums = nums
        self.len = len(nums)

        self.num_idx = {}
        # 移位是为了多加1个"0", 当从0到k个元素是0时, 可以和这个额外添加的"0"匹配
        self.num_idx[0] = 0

        res = []

        prefix_sum = 0
        for idx in range(self.len):
            prefix_sum = prefix_sum + nums[idx]

            # if prefix_sum in self.num_idx:
            same_idx = self.num_idx.get(prefix_sum)
            if same_idx is not None:
                return [same_idx, idx]              # 输出任意一个答案/区间 即可
            else:
                self.num_idx[prefix_sum] = idx+1


def main():
    s = Solution()
    lst = [-3, 1, 2, -3, 4]
    print( s.subarraySum( lst) )

    lst = [-3, 1, -4, 2, -3, 4]
    print( s.subarraySum( lst) )

if __name__ == "__main__":
    main()
